import  os
import requests
from typing import List,Dict,Optional
from chromadb import Client
from dashscope import Generation

# 定义环境变量
CHROMA_DB_PATH= './ragchroma_db'
DASHSCOPE_API_KEY='sk-4a106bdba4ce4b01b836651877585128'
QWEN_MODEL= "qwen-plus"
OLLAMA_URL="http://localhost:11435/api"

# 设置检索和重排参数
TOP_K= 5
RERANK_TOP_K= 3
CHUNK_THRESHOLD= 0.7

class ragdemo:
    def __init__(self):
        # 初始化Chroma DB客户端，使用新的PersistentClient
        from chromadb import PersistentClient, Settings
        self.chroma_client=PersistentClient(
            path=CHROMA_DB_PATH,
            settings=Settings(
                anonymized_telemetry=False
            )
        )
        # 创建或获取集合
        self.collection=self.chroma_client.get_or_create_collection(name="rag_docs")
        print(f"向量数据库已经连接，路径: {CHROMA_DB_PATH}")

        Generation.api_key= DASHSCOPE_API_KEY
        print("通义千问模型已经连接")
    
    # 初步检索
    def _retrieve_from_vector_db(self,query: str)-> List[Dict]:
        results=self.collection.query(
            query_texts=[query],
            n_results=TOP_K
        )
        docs=results['documents'][0]
        metadatas= results['metadatas'][0]
        distances= results['distances'][0]
        return [{
            'content': doc,
            "metadata": meta,
            "distance": float(dist)
        }
        for doc,meta,dist in zip(docs,metadatas,distances)
        if dist<= CHUNK_THRESHOLD
        ]
    
    # 使用bge-m3模型进行重排
    def _rerank_with_bge_m3(self,query: str, candidates: List[Dict])-> List[Dict]:
        # 构造一个重排输入对
        pairs= [[query,item['content']] for item in candidates]
        try:
            response=requests.post(
                f"{OLLAMA_URL}/generate",
                json={
                    "model": "bge-m3",
                    "prompt": f"Rerand these pairs by relevance:\n{pairs}",
                    "options": {"use_embedding_model": True}
                },timeout=10
            )
            print("ollama bge-m3 当前不支持重排，使用原始向量距离排序")
        except:
            print("重排服务不可用")
        ranked=sorted(candidates,key= lambda x: x['distance'])
        return ranked[:RERANK_TOP_K]
    
    # 通过在线通义千问模型回答问题
    def _generate_answer(self,context: str,question: str)-> str:
        prompt = f"""你是一个医疗专家，请根据下面的上下文内容回答问题：
        [上下文]
        {context}
        [问题]
        {question}
        [要求]
        - 如果上下文不足以回答问题，请回复‘根据现有资料无法确认’,
        - 不要编造信息，
        - 使用中文回答
        """
        try:
            response= Generation.Call(
                model=QWEN_MODEL,
                prompt= prompt                
            )
            if response.status_code == 200:
                return response.output.text.strip()
            else:
                return "大模型返回错误"
        except Exception as e:
            print(f"大模型调用失败：{str(e)}")

    # 完整的查询流程：检索，重排，生成答案
    def query(self,user_question: str)-> Dict:
        print(f"\n 用户提问：{user_question}")

        # 初步检索
        raw_results=self._retrieve_from_vector_db(user_question)
        print(f"初步查询到{len(raw_results)}个相关段落")

        if not raw_results:
            return {
                "question": user_question,
                "context":"",
                "answer":"未在知识库中找到相关信息",
                "references":[]
            }
        # 重排
        final_results=self._rerank_with_bge_m3(user_question,raw_results)
        print(f"重排后保留{len(final_results)}个相关段落")

        # 拼接上下文
        context_str= "\n\n".join([item['content'] for item in final_results])

        # 生成答案
        answer=self._generate_answer(context_str,user_question)

        return {
            "question": user_question,
            "context": context_str,
            "answer": answer,
            "references": [
                {
                    "content": item["content"],
                    "source": item["metadata"].get("source","unknown"),
                    "similarity_score": round(1- item["distance"],4)
                }
                for item in final_results
            ]
        }
    
def demo():
    rag=ragdemo()

    test_questions=["高血压的诊断标准是什么？","有个人经常头痛，眩晕，是不是高血压?"]
    for q in test_questions:
        result=rag.query(q)
        print(f"问题：{result['question']}")
        print(f"答案: {result['answer']}")
        print(f"参考来源：")
        for i, ref in enumerate(result["references"]):
            score=ref["similarity_score"]
            src= ref["source"]


demo()

